﻿using System;
using System.Collections.Generic;
using Algoritmia.DataStructures.Maps;
using Algoritmia.Utils;
using Sets = Algoritmia.DataStructures.Sets;
using Algoritmia.DataStructures.Digraphs;

namespace Algoritmia.Problems.SpanningTrees.Prim
{
    public class PrimsMinimumSpanningForestFinder<T> : IMinimumSpanningForestFinder<T,double>, IMinimumSpanningTreeFinder<T,double>
    {
        protected readonly Func<IEnumerable<T>, Sets.ISet<T>> _createSet;
        private readonly Func<IEnumerable<T>, IMap<T, T>> _createFringe;

        public PrimsMinimumSpanningForestFinder(Func<IEnumerable<T>,Sets.ISet<T>> createSet = null,
                                          Func<IEnumerable<T>,IMap<T,T>> createFringe = null
                    )
        {
            _createSet = createSet ?? (it => AlgoritmiaFactory.Sets<T>.Default);
            _createFringe = createFringe ?? (it => AlgoritmiaFactory.Maps<T,T>.Default);
        }

        protected virtual IEnumerable<Tuple<T,T>> MST(IDigraph<T> g, WeightingFunction<T,double> d, T u, Sets.ISet<T> added )
        {
            added.Add(u);
            var inFrontierFrom = _createFringe(g.V);
            foreach(var v in g.Successors(u))
            {
                inFrontierFrom[v] = u;
            }
            while ( inFrontierFrom.Count > 0)
            {
                var optimalEdge = inFrontierFrom.ArgMin(kv => d(kv.Value, kv.Key));
                var v = optimalEdge.Key;
                u = optimalEdge.Value;
                inFrontierFrom.Remove(v);
                added.Add(v);
                yield return new Tuple<T, T>(u,v);
                foreach(var w in g.Successors(v))
                {
                    if (!added.Contains(w) && (!inFrontierFrom.Contains(w) || d(inFrontierFrom[w],w) > d(v,w)))
                    {
                        inFrontierFrom[w] = v;
                    }
                }
            }

        }

        #region Implementation of IMinimumSpanningForestFinder<T,double>

        public IEnumerable<Tuple<T, T>> MinimumSpanningForest(IDigraph<T> g, WeightingFunction<T, double> w)
        {
            var added = _createSet(g.V);
            foreach(var v in g.V)
            {
                if ( !added.Contains(v))
                {
                    foreach(var edge in MST(g,w,v,added))
                    {
                        yield return edge;
                    }
                }
            }
        }

        #endregion

        #region Implementation of IMinimumSpanningTreeFinder<T,double>

        public IEnumerable<Tuple<T, T>> MinimumSpanningTree(IDigraph<T> g, WeightingFunction<T, double> w, T source)
        {
            var added = _createSet(g.V);
            return MST(g, w, source, added);
        }

        #endregion
    }
}
